"""
文档加载和向量化存储模块
负责加载PDF文档，进行文本切分，并存储到Qdrant向量数据库
"""

import os
import logging
from typing import List, Optional
from pathlib import Path

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Qdrant
from langchain_community.embeddings import DashScopeEmbeddings
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams

from utils.config import Config

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class DocumentLoader:
    """文档加载和向量化存储类"""
    
    def __init__(self):
        """初始化文档加载器"""
        # 验证配置
        Config.validate_config()
        
        # 初始化文本分割器
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=Config.CHUNK_SIZE,
            chunk_overlap=Config.CHUNK_OVERLAP,
            length_function=len,
            separators=["\n\n", "\n", "。", "！", "？", "；", " ", ""]
        )
        
        # 初始化嵌入模型
        self.embeddings = DashScopeEmbeddings(
            model=Config.EMBEDDING_MODEL,
            dashscope_api_key=Config.DASHSCOPE_API_KEY
        )
        
        # 初始化Qdrant客户端
        self.qdrant_client = QdrantClient(
            path=Config.QDRANT_PATH
        )
        
        # 确保Qdrant目录存在
        os.makedirs(Config.QDRANT_PATH, exist_ok=True)
        
        logger.info(f"文档加载器初始化完成，Qdrant路径: {Config.QDRANT_PATH}")
    
    def load_pdf_documents(self, data_dir: str) -> List:
        """
        加载指定目录下的所有PDF文档
        
        Args:
            data_dir: 文档目录路径
            
        Returns:
            加载的文档列表
        """
        documents = []
        data_path = Path(data_dir)
        
        if not data_path.exists():
            raise FileNotFoundError(f"数据目录不存在: {data_dir}")
        
        # 查找所有PDF文件
        pdf_files = list(data_path.glob("*.pdf"))
        
        if not pdf_files:
            logger.warning(f"在目录 {data_dir} 中未找到PDF文件")
            return documents
        
        logger.info(f"找到 {len(pdf_files)} 个PDF文件")
        
        # 加载每个PDF文件
        for pdf_file in pdf_files:
            try:
                logger.info(f"正在加载文档: {pdf_file.name}")
                loader = PyPDFLoader(str(pdf_file))
                docs = loader.load()
                
                # 为每个文档添加源文件信息
                for doc in docs:
                    doc.metadata["source"] = pdf_file.name
                    doc.metadata["file_path"] = str(pdf_file)
                
                documents.extend(docs)
                logger.info(f"成功加载文档: {pdf_file.name}, 页数: {len(docs)}")
                
            except Exception as e:
                logger.error(f"加载文档 {pdf_file.name} 时出错: {str(e)}")
                continue
        
        logger.info(f"总共加载了 {len(documents)} 个文档页面")
        return documents
    
    def split_documents(self, documents: List) -> List:
        """
        对文档进行文本切分
        
        Args:
            documents: 原始文档列表
            
        Returns:
            切分后的文档片段列表
        """
        if not documents:
            logger.warning("没有文档需要切分")
            return []
        
        logger.info("开始对文档进行文本切分...")
        split_docs = self.text_splitter.split_documents(documents)
        
        logger.info(f"文档切分完成，共生成 {len(split_docs)} 个文档片段")
        return split_docs
    
    def create_collection(self, collection_name: str = None) -> bool:
        """
        创建Qdrant集合
        
        Args:
            collection_name: 集合名称，默认使用配置中的名称
            
        Returns:
            是否创建成功
        """
        if collection_name is None:
            collection_name = Config.QDRANT_COLLECTION_NAME
        
        try:
            # 检查集合是否已存在
            collections = self.qdrant_client.get_collections()
            existing_collections = [col.name for col in collections.collections]
            
            if collection_name in existing_collections:
                logger.info(f"集合 {collection_name} 已存在")
                return True
            
            # 创建新集合
            self.qdrant_client.create_collection(
                collection_name=collection_name,
                vectors_config=VectorParams(
                    size=1536,  # DashScope text-embedding-v1 的向量维度
                    distance=Distance.COSINE
                )
            )
            
            logger.info(f"成功创建集合: {collection_name}")
            return True
            
        except Exception as e:
            logger.error(f"创建集合时出错: {str(e)}")
            return False
    
    def store_documents_to_vector_db(self, documents: List, collection_name: str = None) -> bool:
        """
        将文档存储到向量数据库
        
        Args:
            documents: 要存储的文档列表
            collection_name: 集合名称
            
        Returns:
            是否存储成功
        """
        if not documents:
            logger.warning("没有文档需要存储")
            return False
        
        if collection_name is None:
            collection_name = Config.QDRANT_COLLECTION_NAME
        
        try:
            # 确保集合存在
            if not self.create_collection(collection_name):
                return False
            
            # 创建向量存储
            vector_store = Qdrant(
                client=self.qdrant_client,
                collection_name=collection_name,
                embeddings=self.embeddings
            )
            
            # 添加文档到向量存储
            logger.info(f"开始将 {len(documents)} 个文档片段存储到向量数据库...")
            vector_store.add_documents(documents)
            
            logger.info("文档存储完成")
            return True
            
        except Exception as e:
            logger.error(f"存储文档到向量数据库时出错: {str(e)}")
            return False
    
    def process_documents(self, data_dir: str, collection_name: str = None) -> bool:
        """
        完整的文档处理流程：加载 -> 切分 -> 存储
        
        Args:
            data_dir: 文档目录路径
            collection_name: 集合名称
            
        Returns:
            是否处理成功
        """
        try:
            # 1. 加载文档
            documents = self.load_pdf_documents(data_dir)
            if not documents:
                logger.error("没有成功加载任何文档")
                return False
            
            # 2. 切分文档
            split_docs = self.split_documents(documents)
            if not split_docs:
                logger.error("文档切分失败")
                return False
            
            # 3. 存储到向量数据库
            success = self.store_documents_to_vector_db(split_docs, collection_name)
            
            if success:
                logger.info("文档处理流程完成")
            else:
                logger.error("文档存储失败")
            
            return success
            
        except Exception as e:
            logger.error(f"文档处理过程中出错: {str(e)}")
            return False
    
    def get_collection_info(self, collection_name: str = None) -> dict:
        """
        获取集合信息
        
        Args:
            collection_name: 集合名称
            
        Returns:
            集合信息字典
        """
        if collection_name is None:
            collection_name = Config.QDRANT_COLLECTION_NAME
        
        try:
            collection_info = self.qdrant_client.get_collection(collection_name)
            return {
                "name": collection_name,
                "vectors_count": collection_info.vectors_count,
                "status": collection_info.status,
                "config": collection_info.config
            }
        except Exception as e:
            logger.error(f"获取集合信息时出错: {str(e)}")
            return {}
